## Pylons DB Admin
## Copyright (C) 2004-2009  Ken Kinder
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
Contains main controller for Pylons DB Admin.
"""
from cStringIO import StringIO
from pylons import request, response, session, tmpl_context as c
from pylons.controllers.util import abort, redirect_to
from pylons.decorators import jsonify
from pylonsdbadmin.lib.base import BaseController, render
from pylonsdbadmin.model.meta import instances
from sqlalchemy.exceptions import SQLAlchemyError
from sqlalchemy.orm.query import Query
import csv
import logging
import math
import simplejson
import sqlalchemy
import sqlalchemy.types
import sqlparse
import time
import urllib

log = logging.getLogger(__name__)

def jsonDefault(o):
    if isinstance(o, sqlalchemy.engine.base.RowProxy):
        return dict(o)
    else:
        return unicode(o)

class ResultContainer(object):
    """
    Each SQL statement/query has a result. Instances of subclasses of this
    class (namely L{ErrorContainer} and L{SuccessContainer}) are used to send
    results to the template in an incapsulated way.
    
    Instances of those subclasses are added to c.results and used in the
    template.
    """
    def isError(self):
        """
        If this ResultContainer instance represents a failed query/statement,
        this method returns true.
        """
        return isinstance(self, ErrorContainer)

class ErrorContainer(ResultContainer):
    """
    See L{ResultContainer} description.
    
    Each instance of this contains information about a query/statement that
    resulted in some kind of SQLAlchemy error (parse error, db error, etc).
    
    @ivar error: Text of error
    @type error: str
    
    @ivar statement: Statement that resulted in this error.
    @type statement: str
    
    @ivar statementParsed: sqlparse Statement object representing statement.
    @type statementParsed: Instance of L{sqlparse.sql.Statement}.
    """
    def __init__(self, error, statement):
        super(ErrorContainer, self).__init__()
        
        self.error = error
        self.statement = statement
        self.statementParsed = sqlparse.parse(statement)[0]

class SuccessContainer(ResultContainer):
    """
    See L{ResultContainer} description.
    
    Each instance of this contains information about a query/statement that
    resulted in success. That success could be actual data, or in the case of
    statements like insert, update, etc, just a number of affected rows.
    
    @ivar resultProxy: SQLAlchemy result proxy that gets data from the
    database. Already iterated over.
    @type resultProxy: Instance of L{SQLAlchemy.engine.ResultProxy}
    
    @ivar data: Iterable (already fetched) with rows and columns from query.
    @type data: list
    
    @ivar statement: Statement that resulted in this error.
    @type statement: str
    
    @ivar statementParsed: sqlparse Statement object representing statement.
    @type statementParsed: Instance of L{sqlparse.sql.Statement}.
    """
    def __init__(self, resultProxy, data, statement):
        super(SuccessContainer, self).__init__()
        
        self.resultProxy = resultProxy
        self.data = data
        self.statement = statement
        self.statementParsed = sqlparse.parse(statement)[0]

class DbadminController(BaseController):
    def index(self):
        c.instances = instances
        return render('/dbadmin/index.mako')
    
    def console(self, instance):
        db = instances[instance]
        db.metadata.reflect()
        
        c.statement = request.params.get('statement', '')
        
        splitStatements = filter(lambda s: s.strip(), sqlparse.split(c.statement))
        
        c.instances = instances
        c.format = c.useFormat = request.params.get('format', 'auto')
        c.instance = instance
        c.db = db
        c.results = []
        
        for i, statement in enumerate(splitStatements):
            #statement = statement.strip()
            s = sqlparse.parse(statement)[0]
            
            if s.token_next_match(0, sqlparse.tokens.Keyword, 'begin'):
                db.session.begin()
            elif s.token_next_match(0, sqlparse.tokens.Keyword, 'rollback'):
                db.session.rollback()
            elif s.token_next_match(0, sqlparse.tokens.Keyword, 'commit'):
                db.session.commit()
            else:
                ##
                ## Data shown in the GUI is auto-limited.
                if c.format in ('auto', 'plain', 'tabs') and s.token_next_match(0, sqlparse.tokens.Keyword, 'select') and not s.token_next_match(0, sqlparse.tokens.Keyword, 'limit'):
                    tokenParts = list(s.flatten())
                    t = tokenParts[-1]
                    
                    if t.ttype == sqlparse.tokens.Punctuation and t.value == ';':
                        if '\n' in statement.strip():
                            tokenParts.insert(-1, '\nlimit 30')
                        else:
                            tokenParts.insert(-1, ' limit 30')
                    else:
                        if '\n' in statement.strip():
                            tokenParts.append('\nlimit 30')
                        else:
                            tokenParts.append(' limit 30')
                    
                    splitStatements[i] = statement = ''.join(unicode(x) for x in tokenParts)
                try:
                    r = db.session.execute(statement)
                except SQLAlchemyError, error:
                    ##
                    ## If we're outputting in a data format and there are
                    ## errors, show the errors instead of writing csv/json
                    ## data.
                    if c.format in ('csv', 'json'):
                        c.useFormat = 'errors-only'
                    c.results.append(ErrorContainer(unicode(error), statement))
                else:
                    assert isinstance(r, sqlalchemy.engine.ResultProxy)
                    if r.closed:
                        c.results.append(SuccessContainer(r, [], statement))
                    else:
                        c.results.append(SuccessContainer(r, r.fetchall(), statement))
                    r.close()
                db.session.flush()
        
        c.statement = ''.join(splitStatements)
        
        if c.format == 'json' and not c.useFormat == 'errors-only':
            response.headers['Content-Type'] = 'application/json'
            data = []
            for result in c.results:
                data.append(result.data)
            return simplejson.dumps(data, default=jsonDefault)
        elif c.format == 'csv' and not c.useFormat == 'errors-only':
            s = StringIO()
            writer = csv.writer(s)
            for resultContainer in c.results:
                if resultContainer.data and getattr(resultContainer.resultProxy, 'keys', None):
                    writer.writerow(resultContainer.resultProxy.keys)
                if resultContainer.data:
                    for row in resultContainer.data:
                        writer.writerow(row)
            
            response.headers['Content-Type'] = 'text/csv'
            response.headers['Content-disposition'] = 'attachment; filename=pylonsdbadmin-out-%s.csv' % (time.strftime('%Y-%m-%d'))
            return s.getvalue()
        else:
            if c.format == 'auto' and len(c.results) > 6:
                c.useFormat = 'plain'
            elif c.format == 'auto' and len(c.results) > 1:
                c.useFormat = 'tabs'
            db.metadata.clear()
            db.metadata.reflect()
            return render('/dbadmin/console.mako')

    def browse(self, instance):
        db = instances[instance]
        db.metadata.reflect()
        c.instance = instance
        c.db = db
        return render('/dbadmin/browse.mako')

    def browse_table(self, instance, table):
        db = instances[instance]
        db.metadata.reflect()
        c.instance = instance
        c.table = table
        c.db = db
        c.dbtable = dbtable = db.metadata.tables[table]
        
        ##
        ## Prepare json-ready colModel, colNames variables required by jqGrid
        colNames = []
        colModel = []
        
        for column in db.metadata.tables[table].columns: 
            colNames.append(column.name)
            m = {'name': column.name}
            colModel.append(m)
        c.colNames = colNames
        c.colModel = colModel
        
        return render('/dbadmin/table.mako')

    @jsonify
    def data(self, instance, table):
        db = instances[instance]
        db.metadata.reflect()
        c.instance = instance
        c.table = table
        c.db = db
        c.dbtable = dbtable = db.metadata.tables[table]
        
        pk = c.dbtable.primary_key.columns.keys()
        
        rows = int(request.params.get('rows', 0))
        page = int(request.params.get('page', 1))
        sidx = request.params.get('sidx', 1)
        sord = request.params.get('sord', 'asc')
        
        q = db.session.query(dbtable)
        count = q.count()
        
        if count > 0 and rows > 0:
            total_pages = math.ceil(float(count)/float(rows))
        else:
            total_pages = 0
        
        if page > total_pages:
            page = total_pages
        
        start = rows*page - rows
        if start < 0:
            start = 0
            
        q = db.session.query(dbtable)
        if request.params.get('sidx', None):
            if request.params.get('sord', u'asc') == 'desc':
                q = q.order_by(request.params['sidx'] + ' desc')
            else:
                q = q.order_by(request.params['sidx'])
        q = q.limit(rows).offset(start)
        data = []
        for row in list(q):
            if len(pk) == 1:
                rowid = getattr(row, pk[0])
            elif len(pk) > 1:
                rowid = []
                for col in pk:
                    rowid.append(urllib.quote_plus(unicode(getattr(row, col))))
                rowid = '/'.join(rowid)
            else:
                return {'error': 'No primary key'}
            data.append({'id': rowid,
                         'cell': list(row)})
        
        return {
            'page': page,
            'total': total_pages,
            'records': count,
            'rows': data}
    
